function [p_policy theta_grid pr1 pr2]= solvemain(ave_theta, true_theta, h_theta, h_eps, delta,cost,wc)
% return discretized theta and policy fun for each point

% add additional country dimension to each output
global num_firm T_sim S T n_theta  beta n_tran tau search mktfe country_num rho_par
% heterogenity of ability
%true_theta = .8; % for simplicity assume true_theta always positive

% choice variable
n = n_tran; % theta discretize #
p = nan(n,(T+1));
sigma = 0.17* cost*exp(ave_theta);
rho = rho_par*sigma;
% discretize points
theta_disc = linspace(ave_theta-2.5*sqrt(1/h_theta)-2.5*sqrt(1/h_eps),ave_theta+2.5*sqrt(1/h_theta)+2.5*sqrt(1/h_eps),n);
% if (true_theta < ave_theta | true_theta == ave_theta)
%     theta_disc = linspace(true_theta-2.5*sqrt(1/h_eps),ave_theta+2.5*sqrt(1/h_eps),n);
% else
%     theta_disc = linspace(ave_theta-2.5*sqrt(1/h_eps),true_theta+2.5*sqrt(1/h_eps),n);
% end;
theta_disc_mat = repmat(theta_disc',1,T+1); %n * T+1

% if (true_theta < ave_theta | true_theta == ave_theta)
%     theta_disc = linspace(true_theta-2.5*sqrt(1/h_theta),ave_theta+2.5*sqrt(1/h_theta),n);
% else
%     theta_disc = linspace(ave_theta-2.5*sqrt(1/h_theta),true_theta+2.5*sqrt(1/h_theta),n);
% end;

% calculate transition matrix
theta_disc_f = [theta_disc(2:n) theta_disc(n)]; %1 by n
theta_disc_l = [theta_disc(1) theta_disc(1:n-1)]; % 1 by n
v = nan(n,n);

h_a_use = wc'; % w_j = -2.41*ave_theta/phi-1, 1 by country_num
theta_disc_f_cmat = repmat(theta_disc_f,1,country_num); % 1 by n * country_num
theta_disc_l_cmat = repmat(theta_disc_l,1,country_num); % 1 by n*country_num
theta_disc_cmat = repmat(theta_disc,1,country_num); % 1 by n*country_num
h_a_use_cmat = reshape(repmat(h_a_use,n,1),1,country_num*n); % 1 by n*country_num
pr1 = normcdf((0.5*(theta_disc_f_cmat+theta_disc_cmat)-(ave_theta*h_theta+true_theta*h_a_use_cmat)./(h_theta+h_a_use_cmat))./(sqrt(h_a_use_cmat)./(h_theta+h_a_use_cmat)))-...
    normcdf((0.5*(theta_disc_l_cmat+theta_disc_cmat)-(ave_theta*h_theta+true_theta*h_a_use_cmat)./(h_theta+h_a_use_cmat))./(sqrt(h_a_use_cmat)./(h_theta+h_a_use_cmat))); 
pr1 = pr1./reshape(repmat(sum(reshape(pr1,n,country_num)),n,1),1,country_num*n); % 1 by n*country_num
% pr(theta_0^k|theta), 1 by n raw probability, when h_a = 1, each row is a country, n*c by 1 probability


h = h_theta +[0:T-1]*h_eps; %1 by T
h_mat = repmat(reshape(repmat(h,n,1),n*T,1),1,n*country_num) + repmat(h_a_use_cmat,n*T,1); %n*T by n*c
% theta_disc
% sprintf('ave:%f, h:%f, true:%f, ha:%f\n',ave_theta, h_theta, true_theta, h_a)

theta_disc_f_mat = repmat(0.5*(theta_disc_f+theta_disc),n*T,1); 
theta_disc_f_mat = repmat(theta_disc_f_mat,1,country_num); %n*T by n*c
theta_disc_l_mat = repmat(0.5*(theta_disc_l+theta_disc),n*T,1); 
theta_disc_l_mat = repmat(theta_disc_l_mat,1,country_num); %n*T by n*c
theta_disc_last_mat = repmat(theta_disc',T,n); 
theta_disc_last_mat = repmat(theta_disc_last_mat,1,country_num);% n*T by n*c
pr2 = normcdf((theta_disc_f_mat-(theta_disc_last_mat.*h_mat+true_theta*h_eps)./(h_mat+h_eps))./(sqrt(h_eps)./(h_eps+h_mat)))-...
    normcdf((theta_disc_l_mat-(theta_disc_last_mat.*h_mat+true_theta*h_eps)./(h_mat+h_eps))./(sqrt(h_eps)./(h_eps+h_mat))); %n*T by n*c
pr2sum = nan(size(pr2));
for i = 1:country_num
    pr2sumtmp = sum(pr2(:,(i-1)*n+1 : i*n),2);
    pr2sum(:,(i-1)*n+1 : i*n)=repmat(pr2sumtmp,1,n);
end
pr2 = pr2./pr2sum;
% pr(theta^k,t|theta^k,t-1) column is all possible state in next period n*T
% by n*c
 
p = nan(n*country_num,T+1); % n*country_num by T+1
mc = cost*exp(true_theta); % net mc of transportation cost
p(:,T+1) = repmat(mc+sigma,n*country_num,1); %n*country_num thetas by T+1 periods
p(:,T+1) = p(:,T+1) + reshape(repmat(tau,1,n)',n*country_num,1); % n*country_num by T+1 
p_use = p(:,T+1);
%in the end
d_mat = exp(1/sigma*(rho*repmat(exp(theta_disc_mat(:,T+1)),country_num,1)-p_use))./reshape(repmat(mktfe,1,n)',n*country_num,1); %n*country_num by 1
d_ind = find(d_mat >1);
d_mat(d_ind) = .99;
%d_mat(d_ind) = 1;
v_final = d_mat*sigma/(1-beta*(1-delta)); %n*county_num by 1 
v = repmat(v_final',n,1); %n by n*county_num
% show probability of purchasing of last period, used to test market size too large or too small
mc_mat = repmat(mc,n*country_num,1); %n*country_num by 1
mc_deliver_mat = mc_mat+reshape(repmat(tau,1,n)',n*country_num,1);
for j = 1:T
    k = T+1-j;
    v_pr = pr2(n*(k-1)+1:n*k,:).*v;
    ev = nan(n*country_num,1); %n*country_num by 1
    ev = sum(v_pr,1)';
    
%     for l = 1:country_num
%         ev(n*(l-1)+1:n*l,:) = sum(v_pr(:,n*(l-1)+1:n*l),2);
%     end
    p_tran = repmat(mc + sigma,n*country_num,1) - beta*(1-delta)*ev + reshape(repmat(tau,1,n)',n*country_num,1); %n*country_num by 1
    indnegative = find(p_tran < mc_deliver_mat );
    p_tran(indnegative) = mc_deliver_mat(indnegative);
    p(:,k) = p_tran;
   % v = repmat(exp(1/sigma*(exp(theta_disc)-p(:,k)'))./((1-beta*(1-delta))*D+beta*(1-delta)*exp(1/sigma*(exp(theta_disc)-p(:,k)')))*sigma,n,1);
    p_use = p_tran ; % n*country_num  by 1
    d_mat = exp(1/sigma.*(rho*repmat(exp(theta_disc_mat(:,k)),country_num,1)-p_use))./reshape(repmat(mktfe,1,n)',n*country_num,1);
    d_ind = find(d_mat >1);
    d_mat(d_ind) = .99;
%    d_mat(d_ind) = 1;
    v = repmat((d_mat*sigma/(1-beta*(1-delta)))',n,1); 

    % v = repmat((sum(d_mat.*repmat(search',n,1),2)*sigma...
    %    /(1-beta*(1-delta)))',n,1); % n by n discre??? adjust buying probability to [0,1]
%     v = repmat((sum(exp(1/sigma.*(repmat(exp(theta_disc)',1,110)-p_use))./repmat(mktfe',n,1).*repmat(search',n,1),2)*sigma)',n,1); % n by n discre

end
%v_a = sum(pr1.*v(1,:));
p_policy = p; %n*country_num by T; price with transportation cost
theta_grid = theta_disc_mat; %possible n state*T+1
end